#!/usr/bin/python

################################################
## A base function for histogram info dumping ##
################################################

#######################################################      
## Set PyRoot Environment and PlotConfig Path first ###
#######################################################
import sys
ROOTSYS = '/afs/atlas.umich.edu/opt/root/lib'
sys.path.append(ROOTSYS)
PlotConfig='../'
PlotConfig2='.'
sys.path.append(PlotConfig)
sys.path.append(PlotConfig2)

#####################
## Import Module  ###
#####################
import array
import os
import glob
from math import sqrt,fabs,sin
from ROOT import TFile,TTree,TChain,TBranch,TH1,TH1F,TList
from ROOT import TLorentzVector,TGraphAsymmErrors,TMath
from ROOT import THStack,TCanvas,TLegend,TColor,TPaveText,TPad
from ROOT import gStyle,gDirectory

# This is a special version to dump histogram information
def dump_info(filetag='',hname='',oname='',configfile='plot_configure',hrange='',hrebin=1,xmax=-1,xmin=-1,verbose=-1,DD_hname='',DD_fname='',DD_process='',**kw):

    '''
    flietag => only include histogram files containing certain tags; 
                in most cases, don't need to set this
                Note that all the histogram should be in ../Result/XXX/
    hname   => the histogram name; to find all the names, simple use 
                "python getlist_histo.py" (loo at histo.list then)
    oname   => the output file name; if not set, will use histogram name
    configfile => Which config file to use, default value is "plot_configure"
    hrange  => A string indicating the ORIGINAL histogram ranges: "nbins,xmin,xmax"; 
                e.g.  "1000,-10,10"
               If you want the program to read the histogram range automatically from input histo, 
               do not set this parameter
    hrebin  => How many times for histogram rebinning
    xmin    => the lower limit for x axis
    xmax    => the upper limit for x axis
    verbose => Set it to be 1, if you want it output verbose message
    '''

    print 'INFO => Dump infomation for histogram %s' % hname
    print ''
    
    ##########################################
    ## Load files,Inputs and Get histograms ##
    ##########################################
    print 'INFO => Loading histograms'
    
    # Parse config file
    from plot_config import plot_config
    dict=plot_config(configfile)
    # Read variables
    IntLumi=dict['IntLumi'] # Integrated Luminosity,pb-1
    MCScale=dict['MCScale'] # Global MCScale Factor
    Dataset=dict['Datasets'] # The Dataset contribution, data/sig/zz/wz...

    # Look at input parameter to see if any Data-Driven backgrounds are required
    list_DD=DD_process.split(',')

    # Get the histogram nbins, xmin and xmax; And declasre histograms for further usage
    histo_range=hrange.split(',')
    histo_all={}
    histo_key={}
    histo_nbin, histo_xmin, histo_xmax = 0, 0., 0.
    if len(histo_range)==3: # if hrange is specified
        histo_nbin, histo_xmin, histo_xmax = int(histo_range[0]),float(histo_range[1]),float(histo_range[2])
    else: # if hrange is not specified, grab a histogram file to extract the histogram range info
        for dataset in Dataset:
            for key in dict[dataset].keys():
                if key=='Color' or key=='Corr' or key=='OverLay': continue
                # get list of folders for this dataset or process
                list_folders=[]
                if 'data' in key: list_folders.append(key)
                else:
                    for folder in dict[dataset][key].keys(): list_folders.append(folder)
                for folder in list_folders:
                    # Read histo file
                    if filetag: filename=glob.glob('%s/*%s*root' % (folder,filetag))
                    else: filename=glob.glob('%s/*root' % (folder))
                    file=TFile(filename[0])
                    histo_temp=gDirectory.Get(hname)
                    if not histo_temp:
                        print 'ERROR ==> In %s, no histogram %s exist!' % (filename,hname)
                        continue
                    histo_nbin=histo_temp.GetNbinsX()
                    histo_xmin=histo_temp.GetXaxis().GetXmin()
                    histo_xmax=histo_temp.GetXaxis().GetXmax()
                    file.Close()
                    break
                break
            break    
    for dataset in Dataset:
        histo_all[dataset]=TH1F(dataset,dataset,histo_nbin, histo_xmin, histo_xmax)
        histo_all[dataset].Sumw2()
    for dataset in Dataset:
        histo_key[dataset]={}
        for key in dict[dataset].keys():
            if key=='Color' or key=='Corr' or key=='OverLay': continue
            histo_key[dataset][key]=TH1F('%s_%s' % (dataset,key),'%s_%s' % (dataset,key),histo_nbin, histo_xmin, histo_xmax)
            histo_key[dataset][key].Sumw2()
    # loop all files and merge histograms (normalization will be done for MC)
    for dataset in Dataset:
        print 'INFO ==> Looking at %s' % dataset
        this_corr=dict[dataset]['Corr']
        for key in dict[dataset].keys():
            if key=='Color' or key=='Corr' or key=='OverLay': continue
            # get list of folders for this dataset or process
            if verbose!=-1:  print 'INFO ==> KEY = %s' % key
            list_folders=[]
            if 'data' in key: list_folders.append(key)
            else: 
                for folder in dict[dataset][key].keys(): list_folders.append(folder)
            # loop all the folders one by one and merge histograms
            xs, k, fe, entry = 0., 0., 0., 0.
            for folder in list_folders:
                # Read histo file
                if filetag: filename=glob.glob('%s/*%s*root' % (folder,filetag))
                else: filename=glob.glob('%s/*root' % (folder))
                file=TFile(filename[0])        
                # Get histogram
                histo_temp=gDirectory.Get(hname)
                if not histo_temp: 
                    print 'ERROR ==> In %s, no histogram %s exist!' % (filename,hname)
                    continue
                # For MC, mormalize histogram and then combine
                if 'mc' in folder:
                    xs=float(dict[dataset][key][folder]['Xsection(pb)'])
                    k=float(dict[dataset][key][folder]['k-factor'])
                    fe=float(dict[dataset][key][folder]['FilterEff(gen)'])
                    entry+=float(dict[dataset][key][folder]['Nevts(Weight)'])
                    if verbose!=-1: print 'INFO ==> Add histograms from %s' % folder
                    histo_key[dataset][key].Add(histo_temp)
                # For Data, just combine histogram inside one dataset
                if 'data' in key:
                    if verbose!=-1: print 'INFO ==> Add histograms from %s' % folder
                    histo_all[dataset].Add(histo_temp)
	            # Close histo file
        	    file.Close()
            # If MC contains several different subsamples in one MCID, one needs to sum 
            # the entries to normalize MC histograms, scale factor = 
            # xs*kfactor*fe*Lumi*MCScale*dict[dataset][Corr]/entry
            if not 'data' in key:
                mc_scale=xs*k*fe*IntLumi*1000.0*MCScale*this_corr/entry
                histo_key[dataset][key].Scale(mc_scale)
                histo_all[dataset].Add(histo_key[dataset][key])
                if verbose!=-1: print 'INFO ==> Normalization done for %s -> %.6f' % (key, mc_scale)

    ####################
    ## Begin to plot ###
    ####################

    # try to get the Data-Driven histograms, and try to do the rebinning and set range automatically
    nDD = len(list_DD)
    list_dd_hname = DD_hname.split(',')
    list_dd_fname = DD_fname.split(',')
    if len(list_dd_hname) !=nDD or len(list_dd_fname)!=nDD :
        print 'ERROR => Data Driven Background have %d processes, not consistent with histo and file lists' % nDD
        sys.exit(-1)
    for i in range(nDD):
        DD_pname = list_DD[i]
        if not DD_pname in Dataset: continue
        print 'INFO ==> Looking at %s Data-Driven' % DD_pname
        DD_fhisto = TFile(list_dd_fname[i])
        DD_histo = gDirectory.Get(list_dd_hname[i])
        if not DD_histo:
            print 'ERROR =>DD Histogram %s not found in root file %d !' % (list_dd_hname[i],list_dd_fname[i])
            sys.exit(-1)
#        DD_nrebin = temp_xwidth/DD_histo.GetBinWidth(1)
#        DD_histo.Rebin(int(DD_nrebin))
        #DD_nbins = DD_histo.GetNbinsX()
        #DD_xmin = DD_histo.GetXaxis().GetXmin()
        #DD_xmax = DD_histo.GetXaxis().GetXmax()
        #A simple treament here, will make it more rebust later
        for j in range(DD_histo.GetNbinsX()):
            ind_bin=histo_all[DD_pname].FindBin(DD_histo.GetBinCenter(j+1))
            histo_all[DD_pname].SetBinContent(ind_bin,DD_histo.GetBinContent(j+1))
            histo_all[DD_pname].SetBinError(ind_bin,DD_histo.GetBinError(j+1))
        DD_fhisto.Close()
        this_corr=dict[DD_pname]['Corr']
        histo_all[DD_pname].Scale(this_corr)

    # rebin
    for dataset in Dataset:
    	histo_all[dataset].Rebin(hrebin)
    # Dump infomation to file
    hmaxb=histo_all['Data'].GetNbinsX()
    if oname=='': output=oname
    else: output=oname
    f = open(output,'w')
    f.write('Information for histogram %s\n' % hname)
    f.write('%10s %10s %10s %10s %10s %10s\n' % ('ndata','nsig','nbkg','edata','esig','ebkg'))
    for bin in range(hmaxb):
        
        ndata,nsig,nbkg,edata,esig,ebkg=0.,0.,0.,0.,0.,0.
        
        ndata = histo_all['Data'].GetBinContent(bin+1)
        edata = histo_all['Data'].GetBinError(bin+1)
        nsig = histo_all['WW->lvlv'].GetBinContent(bin+1)
        esig = histo_all['WW->lvlv'].GetBinError(bin+1)
        for dataset in Dataset:
            if dataset=='Data' or dataset=='WW->lvlv': continue
            else: nbkg += histo_all[dataset].GetBinContent(bin+1)
            ebkg = sqrt( pow(ebkg,2) + pow(histo_all[dataset].GetBinContent(bin+1),2) )
        f.write('%10.4f %10.4f %10.4f %10.4f %10.4f %10.4f\n' % (ndata,nsig,nbkg,edata,esig,ebkg))
    f.close()

    #####################################################
    ## Clear the histograms and canvas after plotting ###
    #####################################################
  
    for dataset in histo_all.keys():
        histo_all[dataset].SetDirectory(0)
        TH1.AddDirectory(0)

if __name__ == "__main__":
        
    dump_info(hname='Njet0_pt2l_ee',oname='info_pt2l_ee',hrebin=10,xmax=2000,xmin=0,configfile='plot_configure_DD_ee',DD_process='W+jets/Dijet',DD_hname='Wjets_CutJetVeto_pt2l_ee',DD_fname='BkDataDriven.root')
    dump_info(hname='Njet0_pt2l_mm',oname='info_pt2l_mm',hrebin=10,xmax=2000,xmin=0,configfile='plot_configure_DD_mm',DD_process='W+jets/Dijet',DD_hname='Wjets_CutJetVeto_pt2l_mm',DD_fname='BkDataDriven.root')
    dump_info(hname='Njet0_pt2l_em',oname='info_pt2l_em',hrebin=10,xmax=2000,xmin=0,configfile='plot_configure_DD_em',DD_process='W+jets/Dijet',DD_hname='Wjets_CutJetVeto_pt2l_em',DD_fname='BkDataDriven.root')
    dump_info(hname='Njet0_pt2l_incl',oname='info_pt2l_incl',hrebin=10,xmax=2000,xmin=0,configfile='plot_configure_DD_incl',DD_process='W+jets/Dijet',DD_hname='Wjets_CutJetVeto_pt2l_incl',DD_fname='BkDataDriven.root')

    dump_info(hname='Njet0_Pt1_ee',oname='info_Pt1_ee',hrebin=10,xmax=2000,xmin=0,configfile='plot_configure_DD_ee',DD_process='W+jets/Dijet',DD_hname='Wjets_CutJetVeto_Pt1_ee',DD_fname='BkDataDriven.root')
    dump_info(hname='Njet0_Pt1_mm',oname='info_Pt1_mm',hrebin=10,xmax=2000,xmin=0,configfile='plot_configure_DD_mm',DD_process='W+jets/Dijet',DD_hname='Wjets_CutJetVeto_Pt1_mm',DD_fname='BkDataDriven.root')
    dump_info(hname='Njet0_Pt1_em',oname='info_Pt1_em',hrebin=10,xmax=2000,xmin=0,configfile='plot_configure_DD_em',DD_process='W+jets/Dijet',DD_hname='Wjets_CutJetVeto_Pt1_em',DD_fname='BkDataDriven.root')
    dump_info(hname='Njet0_Pt1_incl',oname='info_Pt1_incl',hrebin=10,xmax=2000,xmin=0,configfile='plot_configure_DD_incl',DD_process='W+jets/Dijet',DD_hname='Wjets_CutJetVeto_Pt1_incl',DD_fname='BkDataDriven.root')

    dump_info(hname='Njet0_Pt2_ee',oname='info_Pt2_ee',hrebin=10,xmax=2000,xmin=0,configfile='plot_configure_DD_ee',DD_process='W+jets/Dijet',DD_hname='Wjets_CutJetVeto_Pt2_ee',DD_fname='BkDataDriven.root')
    dump_info(hname='Njet0_Pt2_mm',oname='info_Pt2_mm',hrebin=10,xmax=2000,xmin=0,configfile='plot_configure_DD_mm',DD_process='W+jets/Dijet',DD_hname='Wjets_CutJetVeto_Pt2_mm',DD_fname='BkDataDriven.root')
    dump_info(hname='Njet0_Pt2_em',oname='info_Pt2_em',hrebin=10,xmax=2000,xmin=0,configfile='plot_configure_DD_em',DD_process='W+jets/Dijet',DD_hname='Wjets_CutJetVeto_Pt2_em',DD_fname='BkDataDriven.root')
    dump_info(hname='Njet0_Pt2_incl',oname='info_Pt2_incl',hrebin=10,xmax=2000,xmin=0,configfile='plot_configure_DD_incl',DD_process='W+jets/Dijet',DD_hname='Wjets_CutJetVeto_Pt2_incl',DD_fname='BkDataDriven.root')

    dump_info(hname='Njet0_MtWW_ee',oname='info_MtWW_ee',hrebin=10,xmax=2000,xmin=0,configfile='plot_configure_DD_ee',DD_process='W+jets/Dijet',DD_hname='Wjets_CutJetVeto_MtWW_ee',DD_fname='BkDataDriven.root')
    dump_info(hname='Njet0_MtWW_mm',oname='info_MtWW_mm',hrebin=10,xmax=2000,xmin=0,configfile='plot_configure_DD_mm',DD_process='W+jets/Dijet',DD_hname='Wjets_CutJetVeto_MtWW_mm',DD_fname='BkDataDriven.root')
    dump_info(hname='Njet0_MtWW_em',oname='info_MtWW_em',hrebin=10,xmax=2000,xmin=0,configfile='plot_configure_DD_em',DD_process='W+jets/Dijet',DD_hname='Wjets_CutJetVeto_MtWW_em',DD_fname='BkDataDriven.root')
    dump_info(hname='Njet0_MtWW_incl',oname='info_MtWW_incl',hrebin=10,xmax=2000,xmin=0,configfile='plot_configure_DD_incl',DD_process='W+jets/Dijet',DD_hname='Wjets_CutJetVeto_MtWW_incl',DD_fname='BkDataDriven.root')
